from cProfile import label

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator

from LinearRegression.house_price_predict import learning_rate

data = [
    [0.1, 0.3],
    [0.2, 0.35],
    [0.3, 0.41],
    [0.4, 0.48],
    [0.5, 0.54]
]

data = np.asarray(data)

x = data[:, 0]
y = data[:, 1]
x = np.reshape(x, (len(x), 1))
y = np.reshape(y, (len(y), 1))
x = np.hstack([x, np.ones((len(x), 1))])


def hw4_SGD(num_epoch, learning_rate):
    train_losses = []
    theta = np.random.normal(size=x.shape[1])
    theta = np.reshape(theta, (len(theta), 1))
    for i in range(0, num_epoch):
        grad = x.T @ (x @ theta - y)
        theta = theta - learning_rate * grad / len(x)
        train_losses.append([theta[0][0], theta[1][0], np.square(np.mean((x @ theta - y).T @ (x @ theta - y) / 2))])
    return np.asarray(train_losses)


learning_rate = 0.01
num_epoch = 150
train_losses_with_theta = hw4_SGD(num_epoch, learning_rate)
train_losses = train_losses_with_theta[:, 2]
plt.figure()
plt.plot(np.arange(num_epoch), train_losses, color='blue', label='train_losses')
# 由于epoch是整数，这里把图中的横坐标也设置为整数
# 该步骤也可以省略
plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
plt.xlabel('Epoch')
plt.ylabel('RMSE')
plt.legend()
plt.show()
theta_0 = train_losses_with_theta[:, 0]
theta_1 = train_losses_with_theta[:, 1]

# 创建3D图形对象,展示两个维度的参数和损失函数值的关系
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
# 绘制三维散点图
scatter = ax.scatter(theta_0, theta_1, train_losses, c='blue', cmap='viridis')
# 添加颜色条和标签
fig.colorbar(scatter, ax=ax, label='Value')
ax.set_xlabel('theta0')
ax.set_ylabel('theta1')
ax.set_zlabel('train_loss')
ax.set_title('theta-train_losses')

plt.show()